import typing
import numpy as np
import hrvanalysis as hrvana
import matplotlib.pyplot as plt
from sleep_data_obj import SleepData


def get_heart_data_peak_list(raw_data_list: list, sample_rate: int) -> typing.Tuple[list, list]:
    """计算传入的心跳波形的所有峰值点
       前两秒和后两秒的峰值点作为冗余窗口被去除

    Args:
        heart_data_list (list): 心跳波形

    Returns:
        list1: 通用峰值点列表
        list2：求睡眠分期用的峰值点列表
    """
    heart_data = SleepData(sample_rate=sample_rate)
    heart_data.load_data_from_list(raw_data_list)
    raw_data = heart_data.get_data()

    heart_data.remove_mean().filter(
        high_pass_cutoff=1,
        high_pass_filter_order=4,
        high_pass_filter_type='butter',
    )
    raw_data_length = heart_data.get_data_length()

    segment_length = int(10 * sample_rate)
    slide_stride = int(8 * sample_rate)
    # 前1s和后1s为冗余区域，为了消除边界效应，取峰值点时应该将其排除在外
    bias = int(1 * sample_rate)
    max_peak_distance = int(1.5*sample_rate)
    min_peak_distance = int(0.4*sample_rate)
    local_maxium_check_window = int(0.3*sample_rate)
    min_heart_beats = 6

    all_peak_list = []
    for start_index in range(0, raw_data_length, slide_stride):
        end_index = start_index + segment_length
        if end_index > raw_data_length:
            break
        temp_obj = heart_data[start_index:end_index]
        temp_raw_data = raw_data[start_index:end_index]
        # temp_raw_data = temp_obj.get_data()

        part_peak_list = temp_obj.get_peak_list_by_move_avg(
            46, 120,
            rolling_mean_window_size_sec=0.75, edge_length_sec=0.3,
            min_peak_distance_sec=0.4,
            local_maxium_check_window_size_sec=0.1
        )
        part_peak_list = temp_obj.adjust_peaks_position_to_local_maxium(
            temp_raw_data, part_peak_list,
            left_search_points=local_maxium_check_window, right_search_points=local_maxium_check_window
        )

        # exclude peaks which in redundant segment
        part_peak_list = list(
            filter(lambda x: x >= bias and x < slide_stride+bias, part_peak_list))

        if len(part_peak_list) <= min_heart_beats:
            # 不符合最低心率，为噪声片段
            continue

        # 补齐个别缺失的峰值点
        part_peak_list = SleepData.complete_peak_list(
            temp_raw_data, part_peak_list, max_peak_distance, min_peak_distance)
        part_peak_list = SleepData.remove_too_close_peaks(
            part_peak_list, min_peak_distance=int(0.4*sample_rate))
        all_peak_list.append([i+start_index for i in part_peak_list])
    return all_peak_list


def get_hrv(heart_peak_list, sample_rate, mode='HR'):
    rmssd = []
    mean_hr = []
    for i in range(len(heart_peak_list)):
        if(len(heart_peak_list[i]) < 3):
            continue
        temp_heart = np.diff(heart_peak_list[i])
        temp_heart = [j / sample_rate * 1000 for j in temp_heart]
        if mode == 'HR':
            clean_rri = hrvana.remove_outliers(rr_intervals=list(
                temp_heart), low_rri=300, high_rri=2000, verbose=False)  # 300,2000
            clean_rri = hrvana.interpolate_nan_values(
                rr_intervals=clean_rri, interpolation_method="linear")
            clean_rri = hrvana.remove_ectopic_beats(
                rr_intervals=clean_rri, method="malik", verbose=False)
        else:
            clean_rri = hrvana.remove_outliers(rr_intervals=list(
                temp_heart), low_rri=2000, high_rri=6000, verbose=False)  # 300,2000
            clean_rri = hrvana.interpolate_nan_values(
                rr_intervals=clean_rri, interpolation_method="linear")

        temp_heart = hrvana.interpolate_nan_values(rr_intervals=clean_rri)
        heart_time_feature = hrvana.get_time_domain_features(temp_heart)
        # try:
        #     heart_freq_feature = hrvana.get_frequency_domain_features(
        #         temp_heart)
        # except:
        #     continue
        rmssd.append(round(heart_time_feature['rmssd'], 1))
        mean_hr.append(round(heart_time_feature['mean_hr'], 1))
    return rmssd, mean_hr


def _calculate_hrv(raw_ppg_data: list, time_intervals: int, sample_rate: int):
    heart_data_segments_peaks = get_heart_data_peak_list(
        raw_ppg_data, sample_rate)
    if len(heart_data_segments_peaks) < time_intervals/20:
        return -1, -1

    peak_list = []
    for h in heart_data_segments_peaks:
        for p in h:
            peak_list.append(p)

    # 画图
    # plt.plot(raw_ppg_data)
    # plt.scatter(peak_list, [raw_ppg_data[peak] for peak in peak_list], s=35, color='r')
    # plt.show()

    rmssd, mean_hr = get_hrv(heart_data_segments_peaks, sample_rate, "HR")
    # print(np.median(rmssd))
    return np.median(rmssd), np.median(mean_hr)


def calculate_hrv(raw_ppg_data: list, time_intervals: int = 60, sample_rate: int = 25):
    sd = SleepData(sample_rate)
    sd.load_data_from_list(raw_ppg_data)

    hr_results = []
    hrv_results = []
    stride_length = time_intervals*sample_rate
    for start_index in range(0, sd.get_data_length(), stride_length):
        if start_index+stride_length > sd.get_data_length():
            break
        rmssd, hr = _calculate_hrv(
            sd[start_index:start_index+stride_length].get_data(),
            time_intervals, sample_rate
        )
        hrv_results.append(rmssd)
        hr_results.append(hr)
    return hrv_results, hr_results


if __name__ == '__main__':
    hrv_results = calculate_hrv(
        ppg_filepath=r'/Users/lijunyu/Developer/SleepProject/rule_based_dev/ppg_data/putty20220427_02(1).log',
        time_intervals=60,      # 计算间隔是60秒 | 需要改成30秒就将60改为30
        sample_rate=25,
    )
    # 返回值为指定时间间隔的hrv结果列表，-1表示该片段信号质量不好需要重测。
    print(hrv_results)
